'''
本程序仅用于算法部分的实现与基础测试
实际上车版本由C++结合ROS实现

本程序用于模仿节点效果
'''

import cv2
import numpy as np

from rune_detector import *
from rune import *
from tracker import *


#读取视频
# cap = cv2.VideoCapture("D:/PolarBear/#other/能量机关（黑暗环境）/关灯-红方小能量机关-全激活过程.MP4")
cap = cv2.VideoCapture("D:/PolarBear/#other/能量机关（黑暗环境）/关灯-红方大能量机关-正在激活状态.MP4")
# cap = cv2.VideoCapture("D:/PolarBear/#other/能量机关（黑暗环境）/关灯-红方小能量机关-正在激活状态.MP4")
# cap = cv2.VideoCapture("D:/PolarBear/#other/能量机关（黑暗环境）/关灯-红方小能量机关-已激活→不可激活状态.MP4")

# cap = cv2.VideoCapture("D:/PolarBear/#other/能量机关（黑暗环境）/关灯-蓝方大能量机关-全激活过程.MP4")
# cap = cv2.VideoCapture("D:/PolarBear/#other/能量机关（黑暗环境）/关灯-蓝方大能量机关-正在激活状态.MP4")
# cap = cv2.VideoCapture("D:/PolarBear/#other/能量机关（黑暗环境）/关灯-蓝方小能量机关-已激活-静止.MP4")

#创建基本类
class RMRuneDetectorNode():
    def __init__(self):
        # self.detector_ = RuneDetector(thresh=180) # 红色（分离通道版）
        self.detector_ = RuneDetector(thresh=100) # 红色
        # self.detector_ = RuneDetector(thresh=210) # 蓝色

    def ImageCallback(self,image):
        if self.detector_.lose_target_count > 50:
            self.detector_.Reset()
        self.detector_.PreprocessImage(image)
        self.detector_.FindLights()
        self.detector_.DistinguishingLights()
        if self.detector_.state == DectorState.COLLECTING:
            self.detector_.CollectData()
        elif self.detector_.state == DectorState.PREDICTING:
            self.detector_.Predict(0.4)
        
        
node = RMRuneDetectorNode()

#开始读取视频
ret , frame = cap.read()
frame_count = 0
while ret and frame_count<700:
    # print('frame:',frame_count)
    frame_count += 1

    ret , frame = cap.read()
    
    # key = cv2.waitKey(0)
    key = cv2.waitKey(1)
    # key = cv2.waitKey(15)
    # key = cv2.waitKey(30)
    
    if not ret or key == 27:
        break
    if key == 32:
        continue
    
    # if frame_count<240:
    #     continue

    # if len(node.detector_.target_angles_) == 50:
    #     break
    
    node.ImageCallback(frame)
    
    cv2.imshow("binary_img_for_R",node.detector_.binary_img_)
    cv2.imshow("result_img",node.detector_.result_img_)

cap.release()


import matplotlib.pyplot as plt
plt.figure(figsize=(10, 6))# 设置图像尺寸
# 获取到的角位置
x = np.arange(0,len(node.detector_.target_angles_),1)
x = node.detector_.time_stamps_[50:]
y = node.detector_.target_angles_[51:]
plt.plot(x, y, 'g-', label='raw pos')

# 预测的角位置
x = np.array(node.detector_.time_stamps_[50:]) + 0.4
y = node.detector_.predict_angles_
plt.plot(x, y, 'r-', label='predict pos')



# 预测位置与实际位置的差值
x = np.arange(0,len(node.detector_.predict_diffs_),1)
y = node.detector_.predict_diffs_
# plt.plot(x, y, 'r-', label='predict pos diff')

plt.grid(True)  # 添加网格
plt.legend()  # 显示标签
plt.title('Target pos (dt=0.4s)')  # 图形标题
plt.xlabel('time')  # x轴标签
plt.ylabel('pos')  # y轴标签
plt.show()

plt.cla()

# 原始角速度
x = node.detector_.time_stamps_
y = node.detector_.rune_speeds_
plt.plot(x, y, 'g-', label='raw_speed')

# 平滑后的角速度
x = node.detector_.time_stamps_
y = node.detector_.rune_speeds_smoothed_
plt.plot(x, y, 'b-', label='smoothed_speed')

# 拟合后的角速度
params = node.detector_.rune_speed_params
x = node.detector_.time_stamps_
y = np.array([RuneSpeedFunc(x_,params[0],params[1],params[2]) for x_ in x])
plt.plot(x, y, 'r-', label='fitted_speed')

plt.grid(True)  # 添加网格
plt.legend()  # 显示标签
plt.title('Speed')  # 图形标题
plt.xlabel('time')  # x轴标签
plt.ylabel('speed')  # y轴标签
plt.show()